import os.path
import numpy as np
from trsfile import trs_open, Trace, SampleCoding, Header
from trsfile.parametermap import TraceParameterMap, TraceParameterDefinitionMap
from trsfile.traceparameter import ByteArrayParameter, ParameterType, TraceParameterDefinition

def read_header(trs_file_path):
    with trs_open(trs_file_path, 'r') as traces:
        header = {}
        for key, value in traces.get_headers().items():
            header[key.name] = value
        header['PATH']=trs_file_path
        header['data_length']=header['LENGTH_DATA']
        header['trace_points']=header['NUMBER_SAMPLES']
        header['traces_num']=header['NUMBER_TRACES']
        header['file_size']=os.path.getsize(trs_file_path)
        return header

def read_trsfile(trs_file_head, load_traces_once=None, start_pos=None, end_pos=None):
    trs_file_path=trs_file_head['PATH']
    load_traces_once = trs_file_head['NUMBER_TRACES'] if load_traces_once is None else load_traces_once
    start_pos=0 if start_pos is None else start_pos
    end_pos =trs_file_head['NUMBER_TRACES'] if end_pos is None else end_pos
    data_list, traces_list = [], []
    with trs_open(trs_file_path, 'r') as traces:
        for i, trace in enumerate(traces[start_pos:end_pos]):
            data=[]
            for key in trace.parameters:
                data+=trace.parameters[key].value
            data_list.append(data)
            traces_list.append(np.array(trace.samples))
            if (i + 1) % load_traces_once == 0 or (i + 1) == end_pos:
                yield np.array(data_list), np.array(traces_list), (i + 1) == end_pos
                data_list, traces_list = [], []

def write_trsfile(trs_file_path, traces_arr, **kwargs):
    headers = {Header.TRS_VERSION: 2}
    if "description" in kwargs:
        headers[Header.DESCRIPTION]=kwargs['description']
    if "offset_x" in kwargs:
        headers[Header.OFFSET_X] = kwargs['offset_x']
    if "label_x" in kwargs:
        headers[Header.LABEL_X] = kwargs['label_x']
    if "label_y" in kwargs:
        headers[Header.LABEL_Y] = kwargs['label_y']
    if "scale_x" in kwargs:
        headers[Header.SCALE_X] = kwargs['scale_x']
    if "scale_y" in kwargs:
        headers[Header.SCALE_Y] = kwargs['scale_y']
    parameter_def={}
    offset=0
    if "plaintext" in kwargs:
        pt_len=kwargs["plaintext"].shape[1]
        parameter_def["plaintext"]=TraceParameterDefinition(ParameterType.BYTE, pt_len, offset)
        offset+=pt_len
    if "ciphertext" in kwargs:
        ct_len = kwargs["ciphertext"].shape[1]
        parameter_def["ciphertext"] = TraceParameterDefinition(ParameterType.BYTE, ct_len, offset)
        offset += ct_len
    if "mkey" in kwargs and len(kwargs["mkey"].shape) == 2:
        mk_len = kwargs["mkey"].shape[1]
        parameter_def["mkey"] = TraceParameterDefinition(ParameterType.BYTE, mk_len, offset)
        offset += mk_len
    if "mask" in kwargs:
        ms_len = kwargs["mask"].shape[1]
        parameter_def["mask"] = TraceParameterDefinition(ParameterType.BYTE, ms_len, offset)
        offset += ms_len
    headers[Header.LENGTH_DATA]=offset
    headers[Header.TRACE_PARAMETER_DEFINITIONS] = TraceParameterDefinitionMap(parameter_def)
    np_code_dict = {
        "int8": SampleCoding.BYTE, "uint8": SampleCoding.BYTE,
        "int16": SampleCoding.SHORT, "uint16": SampleCoding.SHORT,
        "int32": SampleCoding.INT, "uint32": SampleCoding.INT,
        "int64": SampleCoding.INT, "uint64": SampleCoding.INT,
        "float32": SampleCoding.FLOAT, "float_": SampleCoding.FLOAT, "float64": SampleCoding.FLOAT
    }
    code = np_code_dict[str(traces_arr.dtype)]
    with trs_open(trs_file_path,'w',headers=headers,live_update=True) as traces:
        for i in range(traces_arr.shape[0]):
            trs_param={}
            for key in parameter_def:
                trs_param[key]=ByteArrayParameter(np.array(kwargs[key][i],np.uint8))
            traces.extend(Trace(code,traces_arr[i],TraceParameterMap(trs_param)))